﻿#coding:utf-8
from pydbg import *
from pydbg.defines import *
import utils

#这是程序复位到快照点之后可以单步执行指令的最大次数
MAX_INST = 10

#这个列表不够详细，可以根据自己的要求添加更多的函数
dangerous_functions = {
			"strcpy": "msvcrt.dll",
			"strncpy": "msvcrt.dll",
			"sprintf": "msvcrt.dll",
			"vsprintf": "msvcrt.dll",
			}
dangerous_functions_resolved = {}
crash_encounterd = False
inst_count = 0

def danger_handler(dbg):
	#print the date int the function stack
	#我们将把函数栈上的内容打印出来，通常函数栈上只会有少量参数，
	#我们将提取从ESP到ESP+20的所有值，这应该可提供我们足够的信息来确定
	#这些参数的值是否对用户来说是可控的
	esp_offset = 0
	print "[*] Hit %s" % dangerous_functions_resolved[dbg.context.Eip]
	while esp_offset <= 20:
		parameter = dbg.smart_dereference(dbg.context.Esp+esp_offset)
		print "[Esp + %d] => %s" % (esp_offset, parameter)
		esp_offset += 4
	print "=================================="
	dbg.suspend_all_threads()
	dbg.process_snapshot()
	dbg.resume_all_threads()
	return DBG_CONTINUE

def access_violation_handler(dbg):
	global crash_encounterd
	# We skip first-chance exceptions
	#某些坏事发生了，同时也意味这某些好事发生了
	#处理这个非法访问异常，将进程复位到最近的一次的危险函数调用点
	if dbg.dbg.u.Exception.dwFirstChance:
		return DBG_EXCEPTION_NOT_HANDLED

	crash_bin = utils.crash_binning.crash_binning()
	crash_bin.record_crash(dbg)
	print crash_bin.crash_synopsis()
	if crash_encounterd == False:
		dbg.suspend_all_threads()
		dbg.process_restore()
		crash_encounterd = True
		#我们将每一个线程标志设为单步执行状态
		for thread_id in dbg.enumerate_threads():
			print "[*] Setting single step for thread: 0x%08x" % thread_id
			h_thread = dbg.open_thread(thread_id)
			dbg.single_step(True, h_thread)
			dbg.close_handle(h_thread)
		#现在恢复执行所有线程，这将导致所有控制权马上被
		#single_step_handler处理例程接管
		dbg.resume_all_threads()
		return DBG_CONTINUE
	else:
		dbg.terminate_process()
	return DBG_EXCEPTION_NOT_HANDLED

def single_step_handler(dbg):
	global inst_count
	global crash_encounterd
	if crash_encounterd:
		if inst_count == MAX_INST:
			dbg.single_step(False)
			return DBG_CONTINUE
		else:
			#反汇编下一个单步执行指令
			inst = dbg.disasm(dbg.context.Eip)
			print "#%d\t0x%08x: %s" % (inst_count, dbg.context.Eip, inst)
	return DBG_CONTINUE

pid = raw_input("Enter the Process ID: ")

dbg = pydbg()
dbg.attach(int(pid))

#记录下所有的有风险函数，对其设置断点
for func in dangerous_functions.keys():
	func_address = dbg.func_resolve(dangerous_functions[func],func)
	print "[*] Resolved breakpoint: %s -> 0x%08x" % (func, func_address)
	dbg.bp_set(func_address, handler=danger_handler) # set breakpoint
	dangerous_functions_resolved[func_address] = func

dbg.set_callback(EXCEPTION_ACCESS_VIOLATION,access_violation_handler)
dbg.set_callback(EXCEPTION_SINGLE_STEP, single_step_handler)
dbg.run()